#!/usr/bin/env python3
'''
已适配、未验证

训练记录：
在2号机上训练
20241212：训练报错如下：
49993: done 97 games, mean reward 57.897, speed 222.66 f/s
Traceback (most recent call last):
  File "/home/lutery/projects/python/my_-nqd/learning/atari-carnival/dqn_rainbow.py", line 172, in <module>
    loss_v, sample_prios_v = calc_loss(batch, batch_weights, net, tgt_net.target_model,
  File "/home/lutery/projects/python/my_-nqd/learning/atari-carnival/dqn_rainbow.py", line 60, in calc_loss
    proj_distr = common.distr_projection(next_best_distr, rewards, dones, Vmin, Vmax, N_ATOMS, gamma)
AttributeError: module 'lib.common' has no attribute 'distr_projection'
20241213: 训练分数74分，测试分数103分，继续训练
20241214: 学习率0.0001476225，训练分数79分，测试分数159分，继续训练
20241215：学习率 5.719198113740251e-05，训练分数159分，测试分数238分，继续训练
20241216：学习率2.462e-05,训练分数215分，测试分数285分，继续训练
20241217:学习率 1.4537434250760025e-05，训练分数233分，测试分数359分，继续训练
20241218：学习率6.953209736092344e-06,训练分数256分，测试分数359分，继续训练
20241219: 学习率 4.105800817065169e-06，训练分数260分，测试分数359分，感觉学习率是不是下降的太快了？
20241220：继续训练一天，如果没有下降则play模型看效果，如果效果不好，则重新开始训练并调整学习率下降为4e-06。
学习率2.1819908920219306e-06,训练分数272分，测试分数381分，继续训练
20241221: 学习率： 1.28844380183003e-06，训练分数 278 分，测试分数 381 分，继续训练
20241223：学习率： 8.453479783806829e-07,训练分数283分，测试分数381分，继续训练
20241224：学习率5.546328086155661e-07，训练分数283分，测试分数381分，停止训练，play模型，如果不行，可能是学习率下降的太快
'''
import gymnasium as gym
import ptan
import argparse
import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F

import torch.optim as optim

from tensorboardX import SummaryWriter

from lib import common
from lib.dqn_model import *

import ale_py

gym.register_envs(ale_py)
# n-step
REWARD_STEPS = 4

# priority replay
PRIO_REPLAY_ALPHA = 0.6
BETA_START = 0.4
BETA_FRAMES = 100000



def calc_loss(batch, batch_weights, net, tgt_net, gamma, device="cpu"):
    states, actions, rewards, dones, next_states = common.unpack_batch(batch)
    batch_size = len(batch)

    states_v = torch.tensor(states).to(device)
    actions_v = torch.tensor(actions).to(device)
    next_states_v = torch.tensor(next_states).to(device)
    batch_weights_v = torch.tensor(batch_weights).to(device)

    # next state distribution
    # dueling arch -- actions from main net, distr from tgt_net

    # calc at once both next and cur states
    distr_v, qvals_v = net.both(torch.cat((states_v, next_states_v)))
    next_qvals_v = qvals_v[batch_size:]
    distr_v = distr_v[:batch_size]

    next_actions_v = next_qvals_v.max(1)[1]
    next_distr_v = tgt_net(next_states_v)
    next_best_distr_v = next_distr_v[range(batch_size), next_actions_v.data]
    next_best_distr_v = tgt_net.apply_softmax(next_best_distr_v)
    next_best_distr = next_best_distr_v.data.cpu().numpy()

    dones = dones.astype(np.bool_)

    # project our distribution using Bellman update
    proj_distr = common.distr_projection(next_best_distr, rewards, dones, Vmin, Vmax, N_ATOMS, gamma)

    # calculate net output
    state_action_values = distr_v[range(batch_size), actions_v.data]
    state_log_sm_v = F.log_softmax(state_action_values, dim=1)
    proj_distr_v = torch.tensor(proj_distr).to(device)

    loss_v = -state_log_sm_v * proj_distr_v
    loss_v = batch_weights_v * loss_v.sum(dim=1)
    return loss_v.mean(), loss_v + 1e-5

def test_model(env, net, device, episodes=5):
    with torch.no_grad():
        total_reward = 0.0
        for _ in range(episodes):
            noop_action_count = 0
            pre_action = -1
            obs, _ = env.reset()
            while True:
                obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
                logits_v = net.qvals(obs_v)
                logits = logits_v.data.cpu().numpy()
                action = np.argmax(logits)
                if action == 0 and pre_action == action:  # Noop
                    noop_action_count += 1
                    if noop_action_count > 30:
                        break
                else:
                    noop_action_count = 0
                pre_action = action
                obs, reward, done, trunc, _ = env.step(action)
                total_reward += reward
                if done or trunc:
                    break
    return total_reward / episodes

import os

if __name__ == "__main__":
    params = {
        'env_name':         "ALE/Carnival-v5",
        'stop_reward':      30000,
        'run_name':         'carnival',
        'replay_size': 10 ** 6,
        'replay_initial': 50000,
        'target_net_sync': 10000,
        'epsilon_frames': 10 ** 6,
        'epsilon_start': 1.0,
        'epsilon_final': 0.1,
        'learning_rate': 0.00025,
        'gamma': 0.99,
        'batch_size': 32
    }
    params['epsilon_frames'] *= 2
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    args = parser.parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

    env = common.make_env(params['env_name'])
    test_env = common.make_env(params['env_name'])
    save_path = os.path.join("saves", "rainbow-carnival")
    os.makedirs(save_path, exist_ok=True)

    writer = SummaryWriter(comment="-" + params['run_name'] + "-rainbow")
    net = RainbowDQN(env.observation_space.shape, env.action_space.n).to(device)
    tgt_net = ptan.agent.TargetNet(net)
    print(net)
    agent = ptan.agent.DQNAgent(lambda x: net.qvals(x), ptan.actions.ArgmaxActionSelector(), device=device)

    exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=params['gamma'], steps_count=REWARD_STEPS)
    buffer = ptan.experience.PrioritizedReplayBuffer(exp_source, params['replay_size'], PRIO_REPLAY_ALPHA)
    optimizer = optim.Adam(net.parameters(), lr=params['learning_rate'], eps=1e-3)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=200000, gamma=0.9)

    frame_idx = 0
    train_count = 0
    # 增加加载模型的代码
    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        net.load_state_dict(checkpoint['net'])
        tgt_net.target_model.load_state_dict(checkpoint['tgt_net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        frame_idx = checkpoint['frame_idx']
        train_count = checkpoint['train_count']
        scheduler.load_state_dict(checkpoint['scheduler'])
        print("加载模型成功")
        # 打印学习率
        for param_group in optimizer.param_groups:
            print("学习率：", param_group['lr'])

    beta = BETA_START

    with common.RewardTracker(writer, params['stop_reward']) as reward_tracker:
        while True:
            frame_idx += 1
            buffer.populate(1)
            beta = min(1.0, BETA_START + frame_idx * (1.0 - BETA_START) / BETA_FRAMES)

            new_rewards = exp_source.pop_total_rewards()
            if new_rewards:
                if reward_tracker.reward(new_rewards[0], frame_idx):
                    break

            if len(buffer) < params['replay_initial']:
                continue

            optimizer.zero_grad()
            batch, batch_indices, batch_weights = buffer.sample(params['batch_size'], beta)
            loss_v, sample_prios_v = calc_loss(batch, batch_weights, net, tgt_net.target_model,
                                               params['gamma'] ** REWARD_STEPS, device=device)
            loss_v.backward()
            optimizer.step()
            scheduler.step()
            train_count += 1
            buffer.update_priorities(batch_indices, sample_prios_v.data.cpu().numpy())
            if frame_idx % 8 == 0:
                scheduler.step()

            if frame_idx % params['target_net_sync'] == 0:
                tgt_net.sync()
                checkpoint = {
                    "net": net.state_dict(),
                    "tgt_net": tgt_net.target_model.state_dict(),
                    "optimizer": optimizer.state_dict(),
                    "frame_idx": frame_idx,
                    "train_count": train_count,
                    "scheduler": scheduler.state_dict()
                }
                common.save_checkpoints(train_count, checkpoint, save_path, "rainbow", keep_last=5)
                print(f"Saved checkpoint to {save_path}")

            if frame_idx % 2000 == 0:
                net.eval()
                test_rewards = test_model(test_env, net, device)
                net.train()

                common.save_best_model(test_rewards, net.state_dict(), save_path, "rainbow-best", keep_best=5)
                print(f"save best model, current test score: {test_rewards}")

